package main

import (
	"bytes"
	"flag"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"time"

	"golang.org/x/crypto/ssh"
	"gopkg.in/yaml.v3"
)

// 主机配置结构
type HostConfig struct {
	Host     string   `yaml:"host"`
	Port     int      `yaml:"port,omitempty"`
	User     string   `yaml:"user,omitempty"`
	Password string   `yaml:"password,omitempty"`
	Commands []string `yaml:"commands,omitempty"`
	Su       bool     `yaml:"su,omitempty"`
	SuPass   string   `yaml:"supass,omitempty"`
}

// 公共参数配置
type CommonConfig struct {
	Port     int      `yaml:"port"`
	User     string   `yaml:"user"`
	Password string   `yaml:"password"`
	Commands []string `yaml:"commands,omitempty"`
	Su       bool     `yaml:"su"`
	SuPass   string   `yaml:"supass"`
}

// 配置文件结构
type Config struct {
	Timeout int          `yaml:"timeout"`
	Common  CommonConfig `yaml:"common"`
	Hosts   []HostConfig `yaml:"hosts"`
}

// 声明一个全局变量
var config Config

func main() {
	// 定义命令行参数
	host := flag.String("host", "", "SSH 主机地址")
	port := flag.Int("port", 22, "SSH 端口号")
	user := flag.String("user", "", "SSH 用户名")
	password := flag.String("password", "", "SSH 密码")
	command := flag.String("cmd", "", "要执行的命令")
	timeout := flag.Int("timeout", 10, "连接超时时间(秒)")
	suEnable := flag.Bool("su", false, "是否使用su切换用户")
	// 修改参数名称，从 passwordsu 改为 supass
	suPassword := flag.String("supass", "", "su命令的密码")
	configFile := flag.String("c", "", "配置文件路径或目录路径")
	logFile := flag.String("log", "", "日志文件路径，留空则自动生成")
	// 添加生成默认配置文件的参数
	defaultConfig := flag.Bool("default", false, "生成默认配置文件模板")

	flag.Parse()

	// 如果指定了生成默认配置，则创建配置文件模板
	if *defaultConfig {
		generateDefaultConfig()
		return
	}

	// 生成日志文件名，包含时间戳
	logFilePath := *logFile
	if logFilePath == "" {
		// 生成带有时间戳的日志文件名
		timestamp := time.Now().Format("20060102_150405")
		logFilePath = fmt.Sprintf("ssh_execution_%s.log", timestamp)
	}

	// 获取日志文件的绝对路径
	absLogPath, err := filepath.Abs(logFilePath)
	if err != nil {
		fmt.Printf("无法获取日志文件的绝对路径: %v\n", err)
		os.Exit(1)
	}

	logger, err := os.OpenFile(absLogPath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644)
	if err != nil {
		fmt.Printf("无法打开日志文件: %v\n", err)
		os.Exit(1)
	}
	defer logger.Close()

	// 记录执行开始时间和日志文件位置
	startTime := time.Now().Format("2006-01-02 15:04:05")
	logMessage(logger, fmt.Sprintf("=== 执行开始于 %s ===\n", startTime))
	logMessage(logger, fmt.Sprintf("=== 日志文件: %s ===\n", absLogPath))

	// 如果提供了配置文件或目录，则读取配置
	if *configFile != "" {
		// 检查是文件还是目录
		fileInfo, err := os.Stat(*configFile)
		if err != nil {
			fmt.Printf("无法访问指定的路径: %v\n", err)
			os.Exit(1)
		}

		if fileInfo.IsDir() {
			// 如果是目录，读取所有.yml文件
			processConfigDirectory(*configFile, *timeout, logger)
		} else {
			// 如果是单个文件，直接处理
			processConfigFile(*configFile, *timeout, logger)
		}
	} else {
		// 检查必要参数
		if *host == "" || *user == "" || *password == "" || *command == "" {
			fmt.Println("错误: 主机地址、用户名、密码和命令都是必须的")
			flag.Usage()
			os.Exit(1)
		}

		// 如果启用su但没有提供密码，报错
		if *suEnable && *suPassword == "" {
			fmt.Println("错误: 启用su切换用户时，必须提供su密码")
			flag.Usage()
			os.Exit(1)
		}

		// 使用命令行参数创建单个主机配置
		hostConfig := HostConfig{
			Host:     *host,
			Port:     *port,
			User:     *user,
			Password: *password,
			Commands: []string{*command},
			Su:       *suEnable,
			SuPass:   *suPassword,
		}

		// 创建默认的公共配置
		commonConfig := CommonConfig{
			Port:     22,
			User:     "",
			Password: "",
			Su:       false,
			SuPass:   "",
		}

		// 使用全局变量而不是创建新的局部变量
		config = Config{
			Timeout: *timeout,
			Common:  commonConfig,
			Hosts:   []HostConfig{hostConfig},
		}
	}

	// 设置超时时间，命令行参数优先
	timeoutDuration := time.Duration(config.Timeout) * time.Second
	if *timeout != 10 { // 10是默认值
		timeoutDuration = time.Duration(*timeout) * time.Second
	}

	// 遍历所有主机执行命令
	for _, hostConfig := range config.Hosts {
		// 执行SSH连接和命令
		executeSSHCommands(hostConfig, config.Common, timeoutDuration, logger)
	}

	// 记录执行结束时间
	endTime := time.Now().Format("2006-01-02 15:04:05")
	logMessage(logger, fmt.Sprintf("=== 执行结束于 %s ===\n\n", endTime))
}

// 处理配置目录中的所有.yml文件
func processConfigDirectory(dirPath string, timeoutOverride int, logger *os.File) {
	// 读取目录中的所有文件
	files, err := os.ReadDir(dirPath)
	if err != nil {
		fmt.Printf("无法读取目录: %v\n", err)
		os.Exit(1)
	}

	// 计数器，用于跟踪处理的配置文件数量
	configCount := 0

	// 遍历所有文件
	for _, file := range files {
		// 跳过目录
		if file.IsDir() {
			continue
		}

		// 只处理.yml和.yaml文件
		fileName := file.Name()
		if !strings.HasSuffix(fileName, ".yml") && !strings.HasSuffix(fileName, ".yaml") {
			continue
		}

		// 构建完整的文件路径
		filePath := filepath.Join(dirPath, fileName)

		// 处理配置文件
		logMessage(logger, fmt.Sprintf("\n=== 处理配置文件: %s ===\n", filePath))
		processConfigFile(filePath, timeoutOverride, logger)
		configCount++
	}

	if configCount == 0 {
		fmt.Printf("警告: 在目录 %s 中没有找到任何.yml或.yaml配置文件\n", dirPath)
	} else {
		fmt.Printf("成功处理了 %d 个配置文件\n", configCount)
	}
}

// 处理单个配置文件
func processConfigFile(filePath string, timeoutOverride int, logger *os.File) {
	// 读取配置文件
	data, err := os.ReadFile(filePath)
	if err != nil {
		logMessage(logger, fmt.Sprintf("无法读取配置文件 %s: %v\n", filePath, err))
		return
	}

	var config Config
	err = yaml.Unmarshal(data, &config)
	if err != nil {
		logMessage(logger, fmt.Sprintf("解析配置文件 %s 失败: %v\n", filePath, err))
		return
	}

	// 如果没有主机配置，报错
	if len(config.Hosts) == 0 {
		logMessage(logger, fmt.Sprintf("错误: 配置文件 %s 中没有主机配置\n", filePath))
		return
	}

	// 设置超时时间，命令行参数优先
	timeoutDuration := time.Duration(config.Timeout) * time.Second
	if timeoutOverride != 10 { // 10是默认值
		timeoutDuration = time.Duration(timeoutOverride) * time.Second
	}

	// 遍历所有主机执行命令
	for _, hostConfig := range config.Hosts {
		// 执行SSH连接和命令
		executeSSHCommands(hostConfig, config.Common, timeoutDuration, logger)
	}
}

// 记录日志的辅助函数
func logMessage(logger *os.File, message string) {
	fmt.Print(message)
	logger.WriteString(message)
}

// 执行SSH命令的函数
func executeSSHCommands(hostConfig HostConfig, commonConfig CommonConfig, timeout time.Duration, logger *os.File) {
	// 合并主机配置和公共配置
	port := hostConfig.Port
	if port == 0 {
		port = commonConfig.Port
	}

	user := hostConfig.User
	if user == "" {
		user = commonConfig.User
	}

	password := hostConfig.Password
	if password == "" {
		password = commonConfig.Password
	}

	// 修复：从主机配置或公共配置获取su标志
	su := hostConfig.Su
	// 如果主机配置中没有明确设置su（默认值为false），则使用公共配置中的su值
	if !su {
		su = commonConfig.Su
	}

	suPass := hostConfig.SuPass
	if suPass == "" && su {
		suPass = commonConfig.SuPass
	}

	// 配置 SSH 客户端
	config := &ssh.ClientConfig{
		User: user,
		Auth: []ssh.AuthMethod{
			ssh.Password(password),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		Timeout:         timeout,
	}

	// 连接到 SSH 服务器
	addr := fmt.Sprintf("%s:%d", hostConfig.Host, port)
	logMessage(logger, fmt.Sprintf("\n[%s] 正在连接到 %s...\n", hostConfig.Host, addr))

	client, err := ssh.Dial("tcp", addr, config)
	if err != nil {
		logMessage(logger, fmt.Sprintf("[%s] 无法连接到 SSH 服务器: %v\n", hostConfig.Host, err))
		return
	}
	defer client.Close()

	// 确定要执行的命令
	var commands []string
	if len(hostConfig.Commands) > 0 {
		commands = hostConfig.Commands
	} else {
		commands = commonConfig.Commands
	}

	if len(commands) == 0 {
		logMessage(logger, fmt.Sprintf("[%s] 警告: 没有可执行的命令\n", hostConfig.Host))
		return
	}

	// 执行每个命令
	for _, cmd := range commands {
		// 创建会话
		session, err := client.NewSession()
		if err != nil {
			logMessage(logger, fmt.Sprintf("[%s] 无法创建 SSH 会话: %v\n", hostConfig.Host, err))
			return
		}

		// 设置标准输出和标准错误的缓冲区
		var stdout, stderr bytes.Buffer
		session.Stdout = &stdout
		session.Stderr = &stderr

		logMessage(logger, fmt.Sprintf("[%s] 执行命令: %s\n", hostConfig.Host, cmd))

		// 执行命令
		if su {
			// 使用su切换用户后执行命令
			suCmd := fmt.Sprintf("echo %s | su - -c \"%s\"", suPass, cmd)
			err = session.Run(suCmd)
		} else {
			// 直接执行命令
			err = session.Run(cmd)
		}

		if err != nil {
			logMessage(logger, fmt.Sprintf("[%s] 命令执行失败: %v\n", hostConfig.Host, err))
			logMessage(logger, fmt.Sprintf("[%s] 错误输出:\n%s\n", hostConfig.Host, stderr.String()))
		} else {
			logMessage(logger, fmt.Sprintf("[%s] 命令执行成功，输出:\n%s\n", hostConfig.Host, stdout.String()))
		}

		session.Close()
	}
}

// 生成默认配置文件模板
func generateDefaultConfig() {
	defaultConfig := Config{
		Timeout: 10,
		Common: CommonConfig{
			Port:     22,
			User:     "root",
			Password: "password",
			Commands: []string{"echo 'Hello World'", "ls -la"},
			Su:       false,
			SuPass:   "",
		},
		Hosts: []HostConfig{
			{
				Host:     "192.168.1.100",
				Port:     0,
				User:     "",
				Password: "",
				Commands: []string{},
				Su:       false,
				SuPass:   "",
			},
			{
				Host:     "192.168.1.101",
				Port:     2222,
				User:     "admin",
				Password: "admin123",
				Commands: []string{"df -h", "free -m"},
				Su:       true,
				SuPass:   "rootpass",
			},
		},
	}

	// 将配置转换为YAML
	yamlData, err := yaml.Marshal(&defaultConfig)
	if err != nil {
		fmt.Printf("无法生成YAML: %v\n", err)
		os.Exit(1)
	}

	// 写入文件
	err = os.WriteFile("config.yml", yamlData, 0644)
	if err != nil {
		fmt.Printf("无法写入配置文件: %v\n", err)
		os.Exit(1)
	}

	fmt.Println("默认配置文件 config.yml 已生成")
	fmt.Println("配置文件内容:")
	fmt.Println(string(yamlData))
}
